#include "Stdafx.h"

#include "dmalgointerfaces.h"
#include "dmerrorhandling.h"
#include "DMUtils.h"
#include "dmalgorithm.h"
#include "DMMetadataAndFactory.h"
using namespace Microsoft::SqlServer::DataMining::PluginAlgorithms;

static	const wchar_t*	g_szModelStats	=	L"ModelStats";
static	const size_t	g_cchModelStats	=	wcslen(g_szModelStats);

void* operator new(size_t	sizet, IDMMemoryAllocator& galloc)
{
	byte* pbyte;
	HRESULT hr = galloc.Alloc((UINT) sizet, &pbyte);

	if (!SUCCEEDED(hr))
	{
		return NULL;
	}

	return (void*) pbyte;
}

void operator delete(void* pv, IDMMemoryAllocator& galloc)
{
	galloc.Free((byte*) pv);
}


bool	PMMLSAXContentHandler::ShouldServerParseStatistics()
{
	Object^			obj		=	m_algo;
	AlgorithmBase^	algo	=	dynamic_cast<AlgorithmBase^>(obj);
	return (algo->Model->Metadata->GetMarginalRequirements() != MarginalRequirements::None);
}


System::Xml::XmlWriter^	PMMLSAXContentHandler::Writer()
{
	Object^			obj		=	m_algo;
	AlgorithmBase^	algo	=	dynamic_cast<AlgorithmBase^>(obj);
	return algo->GetPMMLWriter();
}


IDMModelServices*	PMMLSAXContentHandler::ModelServices()
{
	Object^			obj		=	m_algo;
	AlgorithmBase^	algo	=	dynamic_cast<AlgorithmBase^>(obj);
	return algo->Model->m_wpSvc;
}


void PMMLSAXContentHandler::SwitchToServerParsing()
{
	USES_ERROR_HANDLING;
	
	Object^			obj		=	m_algo;
	AlgorithmBase^	algo	=	dynamic_cast<AlgorithmBase^>(obj);
	// Save the depth
	algo->SavePMMLParsingState( m_depth - 1);
	
	THROW_ON_FAILURE( ModelServices()->ParsePMMLModelStatistics( m_wpContext ) );
}

HRESULT PMMLSAXContentHandler::startElement( 
	/* [in] */ const wchar_t *pwchNamespaceUri,
	/* [in] */ int cchNamespaceUri,
	/* [in] */ const wchar_t *pwchLocalName,
	/* [in] */ int cchLocalName,
	/* [in] */ const wchar_t *pwchQName,
	/* [in] */ int cchQName,
	/* [in] */ ISAXAttributes *pAttributes)
{
	USES_ERROR_HANDLING;
	DM_TRY
	{
		m_depth++;

		String^ localName	=	gcnew String(pwchLocalName, 0, cchLocalName);
		String^	ns			=	gcnew String(pwchNamespaceUri, 0, cchNamespaceUri);
		String^	qname		=	gcnew String(pwchQName, 0, cchQName);
		String^	prefix		=	String::Empty;

		if( qname->Length > (localName->Length + 1))
		{
			prefix = qname->Substring( 0, qname->Length - localName->Length - 1);
		}

		if( ShouldServerParseStatistics() )
		{
			// Detect the ModelStats tag and delegate parsing to the server
			if( g_cchModelStats == cchLocalName )
			{
				if( 0 == wcsnicmp( g_szModelStats, pwchLocalName, g_cchModelStats) )
				{
					SwitchToServerParsing();
					REPORT_STATUS;
				}
			}
		}

		

		Writer()->WriteStartElement(prefix, localName, ns);

		

		if( pAttributes != NULL )
		{
			int nLength = 0;
			THROW_ON_FAILURE( pAttributes->getLength(&nLength) );
			for( int nIndex = 0; nIndex < nLength; nIndex ++ )
			{
				const wchar_t*	pwchAttLocalName = NULL;
				int			cchAttLocalName	=	0;
				const wchar_t*	pwchAttQName = NULL;
				int			cchAttQName	=	0;
				const wchar_t*	pwchAttUri = NULL;
				int			cchAttUri	=	0;
				const wchar_t*	pwchValue = NULL;
				int			cchValue	=	0;

				THROW_ON_FAILURE( pAttributes->getLocalName( nIndex, &pwchAttLocalName, &cchAttLocalName) );
				THROW_ON_FAILURE( pAttributes->getQName( nIndex, &pwchAttQName, &cchAttQName) );
				THROW_ON_FAILURE( pAttributes->getURI( nIndex, &pwchAttUri, &cchAttUri) );
				THROW_ON_FAILURE( pAttributes->getValue( nIndex, &pwchValue, &cchValue) );

				localName	=	gcnew String(pwchAttLocalName, 0, cchAttLocalName);
				ns			=	gcnew String(pwchAttUri, 0, cchAttUri);
				qname		=	gcnew String(pwchAttQName, 0, cchAttQName);
				prefix		=	String::Empty;
				String^	value	=	gcnew String( pwchValue, 0, cchValue );

				if( qname->Length > (localName->Length + 1))
				{
					prefix = qname->Substring( 0, qname->Length - localName->Length - 1);
				}

				Writer()->WriteAttributeString(prefix, localName, ns, value);

			}
		}
	}
	DM_CATCHALL;

	REPORT_STATUS;
}

HRESULT PMMLSAXContentHandler::endElement( 
	/* [in] */ const wchar_t *pwchNamespaceUri,
	/* [in] */ int cchNamespaceUri,
	/* [in] */ const wchar_t *pwchLocalName,
	/* [in] */ int cchLocalName,
	/* [in] */ const wchar_t *pwchQName,
	/* [in] */ int cchQName)
{
	USES_ERROR_HANDLING;
	DM_TRY
	{
		Writer()->WriteEndElement();
		m_depth --;
		if( m_depth == 0 )
		{
			THROW_ON_FAILURE( ModelServices()->ContinuePMMLParsing(m_wpContext) );
		}
	}
	DM_CATCHALL;
	REPORT_STATUS;

}

HRESULT PMMLSAXContentHandler::characters( 
	/* [in] */ const wchar_t *pwchChars,
	/* [in] */ int cchChars)
{
	USES_ERROR_HANDLING;
	DM_TRY
	{
		String^ chars	=	gcnew String(pwchChars, 0, cchChars);
		Writer()->WriteString(chars);
	}
	DM_CATCHALL;
	REPORT_STATUS;
}

HRESULT PMMLSAXContentHandler::ignorableWhitespace( 
	/* [in] */ const wchar_t *pwchChars,
	/* [in] */ int cchChars)
{
	USES_ERROR_HANDLING;
	DM_TRY
	{
		String^ chars	=	gcnew String(pwchChars, 0, cchChars);
		Writer()->WriteWhitespace(chars);
	}
	DM_CATCHALL;
	REPORT_STATUS;
}







// PMMLWriter -- class used for serializing PMML
PMMLWriter::PMMLWriter	(ISAXContentHandler* in_pSAXHandler, IMXAttributes* in_pMXAtt)
{
	m_wpContentHandler = in_pSAXHandler;
	m_wpMXAttributes = in_pMXAtt;
	currLevel = -1;
	stackElements = gcnew System::Collections::Stack();
	arTopAttributes = gcnew System::Collections::ArrayList();
	writeState = System::Xml::WriteState::Content;

}
	



void PMMLWriter::Close ()
{
}

void PMMLWriter::Flush ()
{
	CheckState();
}

void PMMLWriter::WriteStartElement (String^ prefix, String^ localName, String^ ns)
{
	CheckState();

	try
	{
		// Move to content
		MoveToContent();

		if( localName == nullptr || localName->Length == 0 )
			throw gcnew System::ArgumentException("localName");


		PMMLElement^ elem = gcnew PMMLElement();
		elem->name = localName;
		if( prefix == nullptr || prefix->Length == 0 )
		{
			elem->qname = localName;
		}
		else
		{
			elem->qname = prefix;
			elem->qname += ":";
			elem->qname += localName;
		}
		elem->ns = (ns == nullptr) ? String::Empty : ns;
		elem->content = String::Empty;

		stackElements->Push( elem );
		currLevel++;

		writeState = System::Xml::WriteState::Element;
	}
	catch(...)
	{
		writeState = System::Xml::WriteState::Error;
		throw;
	}
}

void PMMLWriter::WriteEndElement ()
{
	CheckState();

	try
	{
		if( currLevel < 0 )
		{
			throw gcnew System::InvalidOperationException();
		}

		// Move to content
		MoveToContent();

		// Ends the top element
		EndTopElement();

		currLevel--;
	}
	catch(...)
	{
		writeState = System::Xml::WriteState::Error;
		throw;
	}

}

void PMMLWriter::WriteStartAttribute (String^ prefix, String^ localName, String^ ns)
{
	CheckState();

	try
	{
		if( writeState != System::Xml::WriteState::Element )
			throw gcnew System::InvalidOperationException(writeState.ToString() );

		if( localName == nullptr || localName->Length == 0 )
			throw gcnew System::ArgumentException("localName");

		if( currLevel < 0 )
		{
			throw gcnew System::InvalidOperationException();
		}


		PMMLAttribute^ att = gcnew PMMLAttribute();
		att->name = localName;
		if( prefix == nullptr || prefix->Length == 0 )
		{
			att->qname = localName;
		}
		else
		{
			att->qname = prefix;
			att->qname += ":";
			att->qname += localName;
		}
		att->ns = (ns == nullptr) ? String::Empty : ns;
		att->content = String::Empty;

		arTopAttributes->Add( att );
		writeState = System::Xml::WriteState::Attribute;
	}
	catch(...)
	{
		writeState = System::Xml::WriteState::Error;
		throw;
	}

}

void PMMLWriter::WriteEndAttribute ()
{
	CheckState();

	try
	{
		if( writeState != System::Xml::WriteState::Attribute )
			throw gcnew System::InvalidOperationException(writeState.ToString() );
		if( currLevel < 0 )
		{
			throw gcnew System::InvalidOperationException();
		}


		writeState =  System::Xml::WriteState::Element;
	}
	catch( ...)
	{
		writeState = System::Xml::WriteState::Error;
		throw;
	}
}


void PMMLWriter::WriteChars (array<wchar_t>^ buffer, int index, int count)
{
	CheckState();
	String^ destString = nullptr;

	try
	{
		if( currLevel < 0 )
		{
			throw gcnew System::InvalidOperationException();
		}

		String^ addStr = gcnew System::String(buffer, index, count);
		if( writeState == System::Xml::WriteState::Element || writeState == System::Xml::WriteState::Content )
		{
			
			PMMLElement^	topElem = dynamic_cast<PMMLElement^>(stackElements->Peek());
			topElem->content += addStr;
			MoveToContent();
		}
		else if(writeState == System::Xml::WriteState::Attribute)
		{
			
			PMMLAttribute^	lastAttr = dynamic_cast<PMMLAttribute^>(arTopAttributes[arTopAttributes->Count - 1]);
			lastAttr->content += addStr;
		}
	}
	catch(...)
	{
		writeState = System::Xml::WriteState::Error;
		throw;
	}

}



void PMMLWriter::WriteCData (String^ text)
{
	CheckState();
	USES_ERROR_HANDLING;

	try
	{
		if( currLevel < 0 )
		{
			throw gcnew System::InvalidOperationException();
		}

		CComPtr<ISAXLexicalHandler>		spLexHandler;

		THROW_ON_FAILURE( m_wpContentHandler->QueryInterface( IID_ISAXLexicalHandler, (void**)&spLexHandler ) );
		THROW_ON_FAILURE( spLexHandler->startCDATA() );
		
		BSTR bstrTmp = (BSTR)(void*)System::Runtime::InteropServices::Marshal::StringToBSTR( text );
		HRESULT hRet	=	m_wpContentHandler->characters( bstrTmp,text->Length  );
		System::Runtime::InteropServices::Marshal::FreeBSTR( (IntPtr)(void*)bstrTmp );
		bstrTmp = NULL;

		THROW_ON_FAILURE( hRet );

		THROW_ON_FAILURE( spLexHandler->endCDATA() );
	}
	catch(...)
	{
		writeState = System::Xml::WriteState::Error;
		throw;
	}
}

void PMMLWriter::WriteComment (String^ text)
{
	CheckState();
	USES_ERROR_HANDLING;
	try
	{
		if( currLevel < 0 )
		{
			throw gcnew System::InvalidOperationException();
		}

		CComPtr<ISAXLexicalHandler>		spLexHandler;

		THROW_ON_FAILURE( m_wpContentHandler->QueryInterface( IID_ISAXLexicalHandler, (void**)&spLexHandler) );
		BSTR bstrTmp = (BSTR)(void*)System::Runtime::InteropServices::Marshal::StringToBSTR( text );
		HRESULT hRet	=	spLexHandler->comment(bstrTmp,text->Length  );
		
		System::Runtime::InteropServices::Marshal::FreeBSTR( (IntPtr)(void*)bstrTmp );
		bstrTmp = NULL;

		THROW_ON_FAILURE( hRet );
	}
	catch(...)
	{
		writeState = System::Xml::WriteState::Error;
		throw;
	}

}





void PMMLWriter::WriteFullEndElement ()
{
	CheckState();
	WriteEndElement();
}

void PMMLWriter::WriteWhitespace (String^ ws)
{
	CheckState();
	WriteChars( ws->ToCharArray(), 0, ws->Length);
}

void PMMLWriter::WriteString (String^ text)
{
	CheckState();
	WriteChars( text->ToCharArray(), 0, text->Length);
}
void PMMLWriter::WriteRaw (String^ data)
{
	CheckState();
	WriteChars( data->ToCharArray(), 0, data->Length);
}
void PMMLWriter::WriteRaw (array<wchar_t>^ buffer, int index, int count)
{
	CheckState();
	WriteChars(buffer, index, count);
}







void PMMLWriter::CheckState()
{
	if( this->writeState == System::Xml::WriteState::Error || this->writeState == System::Xml::WriteState::Closed  )
		throw gcnew System::InvalidOperationException();
}


// Methods that are not supported for the PMML Wrter
void PMMLWriter::WriteBase64 (array<unsigned char>^ buffer, int index, int count)
{
	throw gcnew System::NotSupportedException();
}

void PMMLWriter::WriteStartDocument () 
{
	throw gcnew System::NotSupportedException();
}
void PMMLWriter::WriteStartDocument (bool standalone)
{
	throw gcnew System::NotSupportedException();
}

void PMMLWriter::WriteEndDocument ()
{
	throw gcnew System::NotSupportedException();
}


void PMMLWriter::WriteDocType (String^ name, String^ pubid, String^ sysid, String^ subset)
{
	throw gcnew System::NotSupportedException();
}


void PMMLWriter::WriteProcessingInstruction (String^ name, String^ text)
{
	throw gcnew System::NotSupportedException();
}

void PMMLWriter::WriteSurrogateCharEntity (wchar_t lowChar, wchar_t highChar)
{
	throw gcnew System::NotSupportedException();
}

void PMMLWriter::WriteEntityRef (String^ name)
{
	throw gcnew System::NotSupportedException();
}

void PMMLWriter::WriteCharEntity (wchar_t ch)
{
	throw gcnew System::NotSupportedException();
}
String^ PMMLWriter::LookupPrefix (String^ ns)
{
	throw gcnew System::NotSupportedException();
}


void PMMLWriter::MoveToContent()
{
	// Renders startElement and characters for the current element, 
	// resets the content and changes the state to Content
	CheckState();
	USES_ERROR_HANDLING;


	if( writeState == System::Xml::WriteState::Content )
	{
		// In Content already, nothing to do
		return;
	}

	if( writeState != System::Xml::WriteState::Element &&  
		writeState != System::Xml::WriteState::Attribute )
	{
		// Should not be here	
		throw gcnew System::InvalidOperationException(writeState.ToString());
	}
	

	BSTR bstrUri = NULL;
	BSTR bstrLName = NULL;
	BSTR bstrQName = NULL;
	BSTR bstrContent = NULL;


	// const strings
	const wchar_t*		szEmpty = L"";
	const wchar_t*		szCDATA = L"CDATA";

	const wchar_t*	szUri = szEmpty;
	const wchar_t*	szLName = szEmpty;
	const wchar_t*	szQName = szEmpty;
	const wchar_t*	szValue = szEmpty;


	try
	{
		// start, characters then end current element
		m_wpMXAttributes->clear();

		CComPtr<ISAXAttributes> spSAXAtt;
		THROW_ON_FAILURE( m_wpMXAttributes->QueryInterface(&spSAXAtt) );

		// Write attributes into the MXAttributes coclass
		for( int nIndex = 0; nIndex < arTopAttributes->Count; nIndex ++ )
		{
			PMMLAttribute^ attr = dynamic_cast<PMMLAttribute^>(arTopAttributes[nIndex]);

			szUri = szEmpty;
			szLName = szEmpty;
			szQName = szEmpty;
			szValue = szEmpty;

			bstrLName = (BSTR)(void*)System::Runtime::InteropServices::Marshal::StringToBSTR( attr->name );
			szLName	=	bstrLName;

			bstrQName = (BSTR)(void*)System::Runtime::InteropServices::Marshal::StringToBSTR( attr->qname );
			szQName = bstrQName;

			if( attr->ns != nullptr && attr->ns->Length > 0 )
			{
				bstrUri = (BSTR)(void*)System::Runtime::InteropServices::Marshal::StringToBSTR( attr->ns );
				szUri = bstrUri;
			}

			bstrContent = (BSTR)(void*)System::Runtime::InteropServices::Marshal::StringToBSTR( attr->content );
			szValue = bstrContent;

			THROW_ON_FAILURE( m_wpMXAttributes->addAttribute( (BSTR)szUri, (BSTR)szLName, (BSTR)szQName, (BSTR)szCDATA, (BSTR)szValue) ); 

			System::Runtime::InteropServices::Marshal::FreeBSTR( (IntPtr)(void*)bstrUri );
			System::Runtime::InteropServices::Marshal::FreeBSTR( (IntPtr)(void*)bstrLName );
			System::Runtime::InteropServices::Marshal::FreeBSTR( (IntPtr)(void*)bstrQName );
			System::Runtime::InteropServices::Marshal::FreeBSTR( (IntPtr)(void*)bstrContent);
		}
		arTopAttributes->Clear();


		int cUri = 0;
		int cLName = 0;
		int cQName = 0;
		int cContent = 0;

		szUri = szEmpty;
		szLName = szEmpty;
		szQName = szEmpty;
		szValue = szEmpty;


		PMMLElement^ currElement = dynamic_cast<PMMLElement^>(stackElements->Peek());

		bstrLName	=	(BSTR)(void*)System::Runtime::InteropServices::Marshal::StringToBSTR( currElement->name );
		szLName	=	bstrLName;
		cLName	=	currElement->name->Length;
		
		bstrQName	=	(BSTR)(void*)System::Runtime::InteropServices::Marshal::StringToBSTR( currElement->qname );
		szQName	=	bstrQName;
		cQName	=	currElement->qname->Length;

		if( currElement->ns != nullptr && currElement->ns ->Length > 0 )
		{
			bstrUri = (BSTR)(void*)System::Runtime::InteropServices::Marshal::StringToBSTR( currElement->ns );
			szUri = bstrUri;
			cUri	=	currElement->ns->Length;
		}
		
		if( currElement->content != nullptr && currElement->content ->Length > 0 )
		{
			bstrContent = (BSTR)(void*)System::Runtime::InteropServices::Marshal::StringToBSTR( currElement->content );
			szValue = bstrContent;
			cContent = currElement->content->Length;
		}

		THROW_ON_FAILURE( m_wpContentHandler->startElement(szUri, cUri, szLName, cLName, szQName, cQName, spSAXAtt) );
		if( cContent > 0 )
		{
			THROW_ON_FAILURE( m_wpContentHandler->characters(bstrContent, cContent) );
		}

		currElement->content = String::Empty;
		writeState = System::Xml::WriteState::Content;
		
	}
	catch( ... )
	{
		writeState = System::Xml::WriteState::Error;
		throw;
	}
	finally
	{
		if( bstrUri != NULL) 
			System::Runtime::InteropServices::Marshal::FreeBSTR( (IntPtr)(void*)bstrUri );
		if( bstrLName != NULL) 
			System::Runtime::InteropServices::Marshal::FreeBSTR( (IntPtr)(void*)bstrLName );
		if( bstrQName != NULL) 
			System::Runtime::InteropServices::Marshal::FreeBSTR( (IntPtr)(void*)bstrQName );
		if( bstrContent != NULL) 
			System::Runtime::InteropServices::Marshal::FreeBSTR( (IntPtr)(void*)bstrContent );
	}

}


void PMMLWriter::EndTopElement()
{
	// Renders characters for the current element, and changes the state to Content
	CheckState();
	USES_ERROR_HANDLING;

	BSTR bstrUri = NULL;
	BSTR bstrLName = NULL;
	BSTR bstrQName = NULL;
	BSTR bstrContent = NULL;

	try
	{
		if( writeState != System::Xml::WriteState::Content )
		{
			// Should not be here	
			throw gcnew System::InvalidOperationException(writeState.ToString());
		}

		if( currLevel < 0 )
		{
			throw gcnew System::InvalidOperationException();
		}
	



		// const strings
		const wchar_t*		szEmpty = L"";

		const wchar_t*	szUri = szEmpty;
		const wchar_t*	szLName = szEmpty;
		const wchar_t*	szQName = szEmpty;
		const wchar_t*	szValue = szEmpty;


		int cUri = 0;
		int cLName = 0;
		int cQName = 0;
		int cContent = 0;



		PMMLElement^ currElement = dynamic_cast<PMMLElement^>(stackElements->Pop());

		bstrLName	=	(BSTR)(void*)System::Runtime::InteropServices::Marshal::StringToBSTR( currElement->name );
		szLName	=	bstrLName;
		cLName	=	currElement->name->Length;
		
		bstrQName	=	(BSTR)(void*)System::Runtime::InteropServices::Marshal::StringToBSTR( currElement->qname );
		szQName	=	bstrQName;
		cQName	=	currElement->qname->Length;

		if( currElement->ns != nullptr && currElement->ns ->Length > 0 )
		{
			bstrUri = (BSTR)(void*)System::Runtime::InteropServices::Marshal::StringToBSTR( currElement->ns );
			szUri = bstrUri;
			cUri	=	currElement->ns->Length;
		}
		
		if( currElement->content != nullptr && currElement->content ->Length > 0 )
		{
			bstrContent = (BSTR)(void*)System::Runtime::InteropServices::Marshal::StringToBSTR( currElement->content );
			szValue = bstrContent;
			cContent = currElement->content->Length;
		}

		if( cContent > 0 )
		{
			THROW_ON_FAILURE( m_wpContentHandler->characters(bstrContent, cContent) );
		}
		THROW_ON_FAILURE( m_wpContentHandler->endElement(szUri, cUri, szLName, cLName, szQName, cQName) );
		
		writeState = System::Xml::WriteState::Content;
	}
	catch( ...)
	{
		writeState = System::Xml::WriteState::Error;
		throw;
	}
	finally
	{
		if( bstrUri != NULL) 
			System::Runtime::InteropServices::Marshal::FreeBSTR( (IntPtr)(void*)bstrUri );
		if( bstrLName != NULL) 
			System::Runtime::InteropServices::Marshal::FreeBSTR( (IntPtr)(void*)bstrLName );
		if( bstrQName != NULL) 
			System::Runtime::InteropServices::Marshal::FreeBSTR( (IntPtr)(void*)bstrQName );
		if( bstrContent != NULL) 
			System::Runtime::InteropServices::Marshal::FreeBSTR( (IntPtr)(void*)bstrContent );
	}

}